#include "PatternCalibrationProcess.h"
#include <glm/gtx/transform.hpp>


PatternCalibrationProcess::PatternCalibrationProcess () :
		_pattImage (NULL), _pattImageShow (NULL)
{
	int index = 0;
	for (int row = 0; row < PATT_ROWS; ++row) {
		for (int col = 0; col < PATT_COLUMNS; ++col, ++index) {
			_cornersInWorldSpace[index] = glm::vec3 (
					col * PATT_BLOCK_SIZE,
					(PATT_ROWS - row - 1) * PATT_BLOCK_SIZE, 0.0f);
		}
	}

	for (int i = 0; i < PATT_MIN; ++i) {
		glm::vec3 startPoint = _cornersInWorldSpace[i * PATT_COLUMNS + i];
		for (int j = 1; j < PATT_MIN - i; ++j, ++index) {
			_cornersInWorldSpace[index] = startPoint + glm::vec3 (
					0.0f, 0.0f, PATT_BLOCK_SIZE * j);
		}
	}
}

PatternCalibrationProcess::~PatternCalibrationProcess ()
{
}

void PatternCalibrationProcess::initialize ()
{
	_pattImage = cvCreateImage (cvSize (FRAME_WIDTH, FRAME_HEIGHT), 8, 1);
#ifdef SHOW_PATTERN_CALIBRATION_RESULT
	_pattImageShow = cvCreateImage (cvSize (FRAME_WIDTH, FRAME_HEIGHT), 8, 3);
#endif
}

void PatternCalibrationProcess::clean ()
{
	if (_pattImage) cvReleaseImage (&_pattImage);
	if (_pattImageShow) cvReleaseImage (&_pattImageShow);
}

void PatternCalibrationProcess::process (SensorResource &sensorRes)
{
	_pSensor = sensorRes.pSensor;
	_pData = &sensorRes.patternCalibrationData;
	_recognizePattern (sensorRes.pRgbImage);
	glm::mat4x4 rtsmat;
	if (_pData->acc) {
		_calculateCornersInCameraSpace ();
		rtsmat = _calculateExtrinsicMatrixBySolveEquation ();
#ifdef SHOW_PATTERN_CALIBRATION_RESULT
		_showResult (sensorRes.pRgbImage, rtsmat);
#endif
	}

	sensorRes.pSensor->getExtrinsicMat () = rtsmat;
	sensorRes.pSensor->getToWorldMat () = glm::inverse (rtsmat);
}

void PatternCalibrationProcess::_recognizePattern (IplImage *image)
{
	if (!image)
		return;
	cvCvtColor (image, _pattImage, CV_BGR2GRAY);
	CvPoint2D32f corners[PATT_ROWS * PATT_COLUMNS];
	int cornerCount;
	int result = cvFindChessboardCorners (_pattImage,
			cvSize (PATT_COLUMNS, PATT_ROWS), corners, &cornerCount);
	if (result == 1) {
		cvFindCornerSubPix (_pattImage, corners, cornerCount, cvSize (5, 5),
				cvSize (1, 1), cvTermCriteria (CV_TERMCRIT_ITER | CV_TERMCRIT_EPS,
				10, 0.02));
		if (glm::length (glm::vec2 (corners[0].x, corners[0].y)) >
				glm::length (glm::vec2 (corners[PATT_ROWS * PATT_COLUMNS - 1].x,
				corners[PATT_ROWS * PATT_COLUMNS - 1].y))) {
			for (int i = 0; i < PATT_ROWS * PATT_COLUMNS / 2; ++i) {
				CvPoint2D32f temp = corners[i];
				corners[i] = corners[PATT_ROWS * PATT_COLUMNS - 1 - i];
				corners[PATT_ROWS * PATT_COLUMNS - 1 - i] = temp;
			}
		}
		for (int i = 0; i < PATT_ROWS * PATT_COLUMNS; ++i) {
			_pData->cornersAcc[i].x = (_pData->cornersAcc[i].x * _pData->acc + corners[i].x) / (_pData->acc + 1);
			_pData->cornersAcc[i].y = (_pData->cornersAcc[i].y * _pData->acc + corners[i].y) / (_pData->acc + 1);
		}
		++_pData->acc;
	}

}

void PatternCalibrationProcess::_calculateCornersInCameraSpace ()
{
	// glm::vec3 corners3D[PATT_ROWS * PATT_COLUMNS];
	for (int i = 0; i < PATT_ROWS * PATT_COLUMNS; ++i) {
		_cornersInCamSpace[i] = _pSensor->get3DPointInCam (
				_pData->cornersAcc[i].x, _pData->cornersAcc[i].y);
	}
}

void PatternCalibrationProcess::_showResult (IplImage *image, glm::mat4x4 rtsmat)
{
	glm::mat3 perspective = _pSensor->getIntrinsic ().getMatrix ();

	glm::mat4x4 rtsmatR = glm::inverse (rtsmat);
	glm::mat4x3 rtsmat34 (rtsmat[0].swizzle (glm::X, glm::Y, glm::Z),
			rtsmat[1].swizzle (glm::X, glm::Y, glm::Z),
			rtsmat[2].swizzle (glm::X, glm::Y, glm::Z),
			rtsmat[3].swizzle (glm::X, glm::Y, glm::Z));
	glm::vec3 result[PATT_ROWS * PATT_COLUMNS];
	glm::vec4 objPointsInCamera[PATT_ROWS * PATT_COLUMNS];
	glm::vec4 cornersInWorld[PATT_ROWS * PATT_COLUMNS];
	CvMat *result2D = cvCreateMat (1, PATT_ROWS * PATT_COLUMNS, CV_32FC2);

	for (int i = 0; i < PATT_ROWS * PATT_COLUMNS; ++i) {
		objPointsInCamera[i] = rtsmat * glm::vec4 (_cornersInWorldSpace[i].x,
				_cornersInWorldSpace[i].y, _cornersInWorldSpace[i].z, 1);
		cornersInWorld[i] = rtsmatR * glm::vec4 (_cornersInCamSpace[i], 1.0);
		result[i] = perspective * (rtsmat34 * glm::vec4 (
				_cornersInWorldSpace[i].x, _cornersInWorldSpace[i].y,
				_cornersInWorldSpace[i].z, 1));
		result[i] /= result[i].z;
		CV_MAT_ELEM (*result2D, CvPoint2D32f, 0, i) = cvPoint2D32f (
				result[i].x, result[i].y);
	}

	cvCopy (image, _pattImageShow);
	if (_pData->acc)
		cvDrawChessboardCorners (_pattImageShow,
				cvSize (PATT_COLUMNS, PATT_ROWS), _pData->cornersAcc,
				PATT_ROWS * PATT_COLUMNS, 1);
	cvShowImage ((std::string ("Pattern Image - ") +
			_pSensor->getName ()).c_str (), _pattImageShow);

	cvCopy (image, _pattImageShow);
	cvDrawChessboardCorners (_pattImageShow, cvSize (PATT_COLUMNS, PATT_ROWS),
			&CV_MAT_ELEM (*result2D, CvPoint2D32f, 0, 0),
			PATT_ROWS * PATT_COLUMNS, true);
	cvShowImage ((std::string ("New Pattern Image - ") +
			_pSensor->getName ()).c_str (), _pattImageShow);
}

glm::mat4 PatternCalibrationProcess::_calculateExtrinsicMatrix ()
{
	return _calculateTranslation () * _calculateRotation () * _calculateScale ();
}


glm::mat4 PatternCalibrationProcess::_calculateTranslation ()
{
	return glm::translate (_cornersInCamSpace[(PATT_ROWS - 1) * PATT_COLUMNS]);
}

glm::mat4 PatternCalibrationProcess::_calculateRotation ()
{
	glm::vec3 xVec;
	glm::vec3 yVec;
	int xCount = 0;
	int yCount = 0;

	for (int row = 0; row < PATT_ROWS; ++row) {
		for (int col1 = 0; col1 < PATT_COLUMNS; ++col1) {
			glm::vec3 p1 = _cornersInCamSpace[row * PATT_COLUMNS + col1];
			for (int col2 = col1 + 3; col2 < PATT_COLUMNS; ++col2) {
				glm::vec3 p2 = _cornersInCamSpace[row * PATT_COLUMNS + col2];
				glm::vec3 v12 = (p2 - p1) / (float)((col2 - col1) * PATT_BLOCK_SIZE);
				xVec = (xVec * (float)xCount + v12) / (float)(xCount + 1);
				++xCount;
			}
		}
	}

	for (int col = 0; col < PATT_COLUMNS; ++col) {
		for (int row1 = 0; row1 < PATT_ROWS; ++row1) {
			glm::vec3 p1 = _cornersInCamSpace[row1 * PATT_COLUMNS + col];
			for (int row2 = row1 + 3; row2 < PATT_ROWS; ++row2) {
				glm::vec3 p2 = _cornersInCamSpace[row2 * PATT_COLUMNS + col];
				glm::vec3 v21 = (p1 - p2) / (float)((row2 - row1) * PATT_BLOCK_SIZE);
				yVec = (yVec * (float)yCount + v21) / (float)(yCount + 1);
				++yCount;
			}
		}
	}

	glm::vec3 zVec = -glm::normalize (glm::cross (xVec, yVec));
	glm::mat4 rmat (glm::vec4 (xVec, 0.0), glm::vec4 (yVec, 0.0),
			glm::vec4 (zVec, 0.0), glm::vec4 (0.0, 0.0, 0.0, 1.0));
	return rmat;
}

glm::mat4 PatternCalibrationProcess::_calculateScale ()
{
	float _scale = 0.0;
	std::vector< float > scaleArr;
	float minScale = 100;
	float maxScale = 0;
	int count = 0;
	for (int i = 0; i < PATT_ROWS * PATT_COLUMNS; ++i) {
		for (int j = i + 1; j < PATT_ROWS * PATT_COLUMNS; ++j) {
			float distanceWorld = glm::distance (_cornersInWorldSpace[i],
				_cornersInWorldSpace[j]);
			if (distanceWorld < 4) continue;
			float distanceCamera = glm::distance (_cornersInCamSpace[i],
				_cornersInCamSpace[j]);
			float curScale = distanceCamera / distanceWorld;

			_scale = (_scale * count + curScale) / (count + 1);
			++count;

			scaleArr.push_back (curScale);
			if (curScale < minScale)
				minScale = curScale;
			if (curScale > maxScale)
				maxScale = curScale;
		}
	}
	_pData->scale = _scale;
	return glm::scale (1.0f, 1.0f, _scale);
}

glm::mat4 PatternCalibrationProcess::_calculateExtrinsicMatrixCv ()
{
	CvMat objectPoints = cvMat (PATT_ROWS * PATT_COLUMNS, 1, CV_32FC3,
			&_cornersInWorldSpace[0][0]);
	CvMat imagePoints = cvMat (PATT_ROWS * PATT_COLUMNS, 1, CV_32FC2,
			&_pData->cornersAcc[0].x);
	glm::vec3 rVec, tVec;
	glm::mat3 rotateMatrix;
	CvMat rVecM = cvMat (3, 1, CV_32FC1, &rVec[0]);
	CvMat rMat = cvMat (3, 3, CV_32FC1, &rotateMatrix[0][0]);
	CvMat tVecM = cvMat (3, 1, CV_32FC1, &tVec[0]);
	cvFindExtrinsicCameraParams2 (&objectPoints, &imagePoints,
			_pSensor->getIntrinsic ().matrixCvMat (), NULL, &rVecM, &tVecM);
	cvRodrigues2 (&rVecM, &rMat);
	_calculateScale ();
	glm::mat4 rtsMatrix (glm::transpose (rotateMatrix));
	rtsMatrix = glm::translate (tVec) * rtsMatrix;
	rtsMatrix = glm::scale (_pData->scale, _pData->scale, _pData->scale) * rtsMatrix;
	return rtsMatrix;
}

glm::mat4 PatternCalibrationProcess::_calculateExtrinsicMatrixBySolveEquation ()
{
	for (int i = 0, index = PATT_ROWS * PATT_COLUMNS; i < PATT_MIN; ++i) {
		glm::vec3 startPoint = _cornersInCamSpace[i * PATT_COLUMNS + i];
		for (int j = 1; j < PATT_MIN - i; ++j, ++index) {
			glm::vec3 p1 = _cornersInCamSpace[i * PATT_COLUMNS + i + j];
			glm::vec3 v1 = p1 - startPoint;
			glm::vec3 p2 = _cornersInCamSpace[(i + j) * PATT_COLUMNS + i];
			glm::vec3 v2 = p2 - startPoint;
			glm::vec3 vz = glm::cross (v1, v2);
			vz = glm::normalize (vz) * glm::sqrt (glm::length (vz));
			_cornersInCamSpace[index] = startPoint + vz;
		}
	}

	CvMat *pCornersInCamH = cvCreateMat (PATT_ROWS * PATT_COLUMNS + PATT_Z_POINTS, 4, CV_32FC1);
	for (int i = 0; i < PATT_ROWS * PATT_COLUMNS + PATT_Z_POINTS; ++i) {
		CV_MAT_ELEM (*pCornersInCamH, float, i, 0) = _cornersInCamSpace[i].x;
		CV_MAT_ELEM (*pCornersInCamH, float, i, 1) = _cornersInCamSpace[i].y;
		CV_MAT_ELEM (*pCornersInCamH, float, i, 2) = _cornersInCamSpace[i].z;
		CV_MAT_ELEM (*pCornersInCamH, float, i, 3) = 1.0f;
	}
	CvMat cornersInWorld = cvMat (PATT_ROWS * PATT_COLUMNS + PATT_Z_POINTS, 3, CV_32FC1,
			&_cornersInWorldSpace[0][0]);
	CvMat *pTransformMat = cvCreateMat (4, 3, CV_32FC1);
	cvSolve (pCornersInCamH, &cornersInWorld, pTransformMat, CV_SVD);

	glm::mat4 rtsMatrix;
	for (int col = 0; col < 4; ++col)
		for (int row = 0; row < 3; ++row)
			rtsMatrix[col][row] = CV_MAT_ELEM (*pTransformMat, float, col, row);
	rtsMatrix = glm::inverse (rtsMatrix);
	cvReleaseMat (&pCornersInCamH);
	cvReleaseMat (&pTransformMat);
	return rtsMatrix;
}

